package rsa

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"strings"
)

const (
	PriPemBegin = "-----BEGIN PRIVATE KEY-----\n"
	PriPemEnd   = "\n-----END PRIVATE KEY-----"

	PubPemBegin = "-----BEGIN PUBLIC KEY-----\n"
	PubPemEnd   = "\n-----END PUBLIC KEY-----"
)

func parsePrivateKey(privateKey string) (*rsa.PrivateKey, error) {
	privateKey = formatPrivateKey(privateKey)
	// 2、解码私钥字节，生成加密对象
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		return nil, errors.New("私钥信息错误！")
	}
	// 3、解析DER编码的私钥，生成私钥对象
	priKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		if strings.Contains(err.Error(), "use ParsePKCS1PrivateKey instead for this key format") {
			// 尝试解析PKCS1格式的私钥
			priKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
			if err != nil {
				return nil, err
			}

			// 4、断言私钥对象类型
			rsaPriKey, ok := priKey.(*rsa.PrivateKey)
			if !ok {
				return nil, errors.New("私钥类型错误！")
			}

			return rsaPriKey, nil
		}
		return nil, err
	}

	// 4、断言私钥对象类型
	rsaPriKey, ok := priKey.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("私钥类型错误！")
	}

	return rsaPriKey, nil
}

func formatPrivateKey(privateKey string) string {
	if !strings.HasPrefix(privateKey, PriPemBegin) {
		privateKey = PriPemBegin + privateKey
	}
	if !strings.HasSuffix(privateKey, PriPemEnd) {
		privateKey = privateKey + PriPemEnd
	}
	return privateKey
}

func parsePublicKey(publicKey string) (*rsa.PublicKey, error) {
	publicKey = formatPublicKey(publicKey)
	// 2、解码私钥字节，生成加密对象
	block, _ := pem.Decode([]byte(publicKey))
	if block == nil {
		return nil, errors.New("公钥信息错误！")
	}

	// 3、解析DER编码的私钥，生成公钥对象
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		if strings.Contains(err.Error(), "use ParsePKCS1PublicKey instead for this key format") {
			// 解析PKCS1格式的公钥
			pubKey, err = x509.ParsePKCS1PublicKey(block.Bytes)
			if err != nil {
				return nil, err
			}

			// 4、断言公钥对象类型
			rsaPubKey, ok := pubKey.(*rsa.PublicKey)
			if !ok {
				return nil, errors.New("公钥类型错误！")
			}

			return rsaPubKey, nil
		}
		return nil, err
	}

	// 4、断言公钥对象类型
	rsaPubKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("公钥类型错误！")
	}

	return rsaPubKey, nil
}

func formatPublicKey(privateKey string) string {
	if !strings.HasPrefix(privateKey, PubPemBegin) {
		privateKey = PubPemBegin + privateKey
	}
	if !strings.HasSuffix(privateKey, PubPemEnd) {
		privateKey = privateKey + PubPemEnd
	}
	return privateKey
}
